package com.caomeisan.muzi.kcface;

import android.Manifest;
import android.content.pm.PackageManager;
import android.os.Bundle;
import android.text.TextUtils;
import android.util.Log;
import android.view.View;

import androidx.appcompat.app.AppCompatActivity;
import androidx.core.app.ActivityCompat;
import androidx.core.content.ContextCompat;

import com.kuangchi.faiss.utils.FaissJni;

import java.io.BufferedInputStream;
import java.io.BufferedOutputStream;
import java.io.DataInputStream;
import java.io.DataOutputStream;
import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;

public class MainActivityBeiFen extends AppCompatActivity {


    FaissJni faissJni;
    public static final int FEATURE_SIZE = 1024;
    private static List<FaceRegisterInfo> faceRegisterInfoList = new ArrayList<>();
    private static final String path = "/sdcard/SourceInfo/";
    private static final int REQUEST_READ_EXTERNAL_STORAGE = 1;
    private static final String TYPE_FLATIP = "FlatIP";
    private static final String TYPE_IVFFLATIP = "IVFFlatIP";
    private static final String TYPE_IVFPQ = "IVFPQ";
    private static final String TAG = "muzi102";




    @Override
    protected void onCreate(Bundle savedInstanceState) {
        super.onCreate(savedInstanceState);
        setContentView(R.layout.activity_main);
        faissJni = new FaissJni();


        if (ContextCompat.checkSelfPermission(this, Manifest.permission.WRITE_EXTERNAL_STORAGE)
                != PackageManager.PERMISSION_GRANTED) {
            // 如果没有权限，请求权限
            ActivityCompat.requestPermissions(this,
                    new String[]{Manifest.permission.READ_EXTERNAL_STORAGE, Manifest.permission.WRITE_EXTERNAL_STORAGE},
                    REQUEST_READ_EXTERNAL_STORAGE);
        } else {
            // 已经有权限，执行读取操作
            // ...
            //        FaissJni.freeJNI(1000);
            startTesting();
        }


        findViewById(R.id.btn_search).setOnClickListener(new View.OnClickListener() {
            @Override
            public void onClick(View view) {
                if (faceRegisterInfoList != null) {
                    long startTime = System.currentTimeMillis();
                    Log.e(TAG, "开始查询");
                    for (int i = 0; i < faceRegisterInfoList.size(); i++) {
                        FaceRegisterInfo faceRegisterInfo = faceRegisterInfoList.get(i);
                        search(faiss_index, faceRegisterInfo.getFeatureData());
                    }

//                    private void search(int nb, int d, int nq, int topK, long faiss_index, int[] ids, float[] data) {
                    Log.e(TAG, "耗时 = " + (System.currentTimeMillis() - startTime));
//
                }
            }
        });
    }


    void startTesting() {
        File file = new File("/sdcard/faceSourceInfo/");
        if (!file.exists()) {
            boolean b = file.mkdirs();
        }
        File path1 = new File(path);
        if (!path1.exists()) {
            path1.mkdirs();
        }

        File[] featureFiles = file.listFiles();
        if (featureFiles == null || featureFiles.length == 0) {
            return;
        } else {
            long time = System.currentTimeMillis();
            for (File featureFile : featureFiles) {
                try {
                    FileInputStream fis = new FileInputStream(featureFile);
                    byte[] feature = new byte[FEATURE_SIZE];
                    fis.read(feature);
                    faceRegisterInfoList.add(new FaceRegisterInfo(feature, featureFile.getName()));
                    fis.close();
                } catch (IOException e) {
                    e.printStackTrace();
                }
            }
            Log.e(TAG, "时间 = " + (System.currentTimeMillis() - time));
        }
        process(TYPE_IVFFLATIP, 10000);
    }


    private int nb = 20000;
    private int d = 256;
    private int nq = 20000;
    private int topk = 1;



    private long faiss_index = 0;
    float[] mData = null;


    //初始化
    private void process(String indexType, Integer nq) {

//        int nb = 200000;  // 向量集大小
//        int nb = 20000;  // 向量集大小
//        int d = 256;      // 向量维度
//		int nq = 10000;   // 查询次数
        int topK = 1;     // 查找最相似的k个向量
        int[] ids = new int[nb];
        float[] data = null;


        String index_file_path = path + "faces_%s.index";
        String features_file_path = path + "features.dat";

        data = load_features(nb, d, features_file_path, ids);

        index_file_path = String.format(index_file_path, indexType);

        faiss_index = load_index(nb, d, indexType, index_file_path, ids, data);
        mData = data;

    }

    private long load_index(int nb, int d, String type, String index_file_path, int[] ids, float[] data) {
        long t_index = -1;
        File indexFile = new File(index_file_path);
//        logger.info("indexFile:{}", indexFile.getAbsolutePath());
        if (indexFile.exists()) {
//            logger.info("read faiss index");
            t_index = faissJni.readIndexJNI(index_file_path);
        } else {
            File newFile = new File(index_file_path);
            if (!newFile.exists()) {
                try {
                    newFile.createNewFile();
                } catch (IOException e) {
                    e.printStackTrace();
                }
            }
            if ("IVFFlatIP".equals(type)) {
                long nlist = 1024;  // 聚类中心的个数
                long nprobe = 16;   // 查找聚类中心的个数，默认为1个。
                t_index = faissJni.indexIVFFlatJNI(d, nlist, nprobe);
            } else if ("IVFPQ".equals(type)) {
                long nlist = 512;  // 聚类中心的数量
                long m = 64;        // 每个向量被划分为m个子向量
                long nbits = 8;     // 每个子向量使用nbits位进行量化
                t_index = faissJni.indexIVFPQJNI(d, nlist, m, nbits);
            } else {
                t_index = faissJni.indexFlatIPJNI(d);
            }
            // 在一些索引结构（比如IndexIVF及其变种）中，需要在索引之前对向量进行训练，以便索引能够更好地适应数据的分布。
            if ("IVFFlatIP".equals(type) || "IVFPQ".equals(type)) {
//                logger.info("train data");
                faissJni.trainJNI(t_index, nb, data);
            }
//            logger.info("add data");
            // 向索引中添加向量
            faissJni.addJNI(t_index, nb, data);
//            logger.info("write index");
            // 将索引保存到文件
            faissJni.writeIndexJNI(t_index, index_file_path);
        }
        return t_index;
    }


    //加载本地的特征
    private float[] load_features(int nb, int d, String features_file_path, int[] ids) {
        float[] data = null;
        File featuresFile = new File(features_file_path);

        if (featuresFile.exists()) {
            // 从文件中加载id和data
            data = new float[nb * d];
            readFeatures(nb, d, featuresFile, ids, data);
        } else {

            File newFile = new File(features_file_path);
            if (!newFile.exists()) {
                try {
                    newFile.createNewFile();
                } catch (IOException e) {
                    e.printStackTrace();
                }
            }

            data = FaissUtil.makeRandomNormalizedFloatArray(nb, d);

            for (int i = 0; i < nb; i++) {
                ids[i] = i + 100;
            }
            // 保存id和data到文件

            writeFeatures(nb, d, featuresFile, ids, data);
        }
        return data;
    }


    private void readFeatures(int nb, int d, File featuresFile, int[] ids, float[] data) {
        try (DataInputStream inputStream = new DataInputStream(new BufferedInputStream(new FileInputStream(featuresFile)))) {

            for (int i = 0; i < nb; i++) {
                ids[i] = inputStream.readInt();
                for (int j = 0; j < d; j++) {
                    data[i * d + j] = inputStream.readFloat();
                }
            }

//            logger.info("从文件中加载id和data。");
        } catch (IOException e) {
//            logger.error("从文件中加载id和data报错！", e);
        }

    }

    private void writeFeatures(int nb, int d, File featuresFile, int[] ids, float[] data) {


        if (!featuresFile.exists()) {
            try {
                featuresFile.createNewFile();
            } catch (IOException e) {
                e.printStackTrace();
            }
        }

        try (DataOutputStream outputStream = new DataOutputStream(new BufferedOutputStream(new FileOutputStream(featuresFile)))) {
            for (int i = 0; i < nb; i++) {
                outputStream.writeInt(ids[i]);
                for (int j = 0; j < d; j++) {
                    outputStream.writeFloat(data[i * d + j]);
                }
            }
//            logger.info("保存id和data到文件。");
        } catch (IOException e) {
//            logger.error("保存id和data到文件报错！", e);
        }
    }


    //搜索
    private void search(long faiss_index, byte[] features) {

        // 执行nq次查询并计算耗时和topK正确率

        long total_time = 0l;
        int query_vector_count = 1;

        float[] query_vector = ByteUtils.byte2Float(features);
//        Log.e(TAG, "特征float[] = " + Arrays.toString(query_vector));

        //归一化
        float[] query_vectorNew = ByteUtils.normalizeNew01(query_vector);

//        Log.e(TAG, Arrays.toString(query_vector));


//        for (int j = 0; j < d; j++) {
//            query_vector[j] = mData[(int) 99999 * d + j];
//        }

        long start_time = System.currentTimeMillis();
        // 查找最相似的topK个向量
        String resultStr = faissJni.searchJNI(faiss_index, query_vector_count, query_vectorNew, 1);


        long end_time = System.currentTimeMillis();
        total_time = end_time - start_time;
        if (!TextUtils.isEmpty(resultStr)) {
            String[] resultArray = resultStr.split(",");
            String[] topOne = resultArray[0].split(":");
            long query_index = Long.parseLong(topOne[0]);
            Log.e(TAG, " 命中：query_index  = " + query_index + "" + " 命中：目标向量  = " + Float.parseFloat(topOne[1]) + " 命中：相似度 = " + ByteUtils.similarity(Float.parseFloat(topOne[1])));
        } else {
            Log.e(TAG, "未命中：目标向量{},匹配结果为空 ");
        }
        // 计算top1正确率
        Log.e(TAG, "总耗时: 毫秒 = " + total_time);

    }

    @Override
    public void onRequestPermissionsResult(int requestCode, String[] permissions, int[] grantResults) {
        super.onRequestPermissionsResult(requestCode, permissions, grantResults);
        if (requestCode == REQUEST_READ_EXTERNAL_STORAGE) {
            if (grantResults.length > 0 && grantResults[0] == PackageManager.PERMISSION_GRANTED) {
                // 用户授予了权限，可以执行读取操作
                // ...
            } else {
                // 用户拒绝了权限，可能需要提供解释或采取其他措施
                // ...
            }
        }
    }
}